"""
关联规则
一个项集是非频繁项集，那么它的所有超集也是非频繁的。
"""


def load_dataset():
    return [[1, 3, 4], [2, 3, 5], [1, 2, 3, 5], [2, 5]]


def create_c1(dataset):
    """
    初始化候选集
    :param dataset: 数据集
    :return: 候选集
    """
    c1 = []
    for item_set in dataset:
        for item in item_set:
            if [item] not in c1:
                c1.append([item])
    c1.sort()   # 排序
    # frozenset 不可变集合，可以成为字典key，也保留了集合的基本操作
    return list(map(frozenset, c1))


def scan_d(dataset, ck, min_support):
    """
    扫描并筛选满足支持度的项集
    :param dataset: 数据集
    :param ck: 各候选集
    :param min_support: 最小支持度
    :return: 满足支持度的频繁项集, 相应支持度
    """
    # 对每个项集进行计数
    count = {}
    for key in ck:
        for item_set in dataset:
            # 判断 idx 是否是 item_set 的子集
            if key.issubset(item_set):
                if key not in count:
                    count[key] = 1
                else:
                    count[key] += 1
    # 计算支持度，筛选出满足支持度的频繁项集
    num_item = float(len(dataset))
    ret_list = []
    support_data = {}
    for key in count:
        support = count[key] / num_item
        if support >= min_support:
            ret_list.append(key)
            support_data[key] = support
    return ret_list, support_data


def apriori_gen(item_set, k):
    """
    拼接新的候选集
    :param item_set: 项
    :param k: 判断是否拼接，提高效率
    :return: 拼接好的候选集
    """
    ret_list = []
    n = len(item_set)
    for i in range(n):
        for j in range(i + 1, n):
            # 除最后一项外，前项全部一样的进行拼接
            item1 = list(item_set[i])[:k - 2]
            item2 = list(item_set[j])[:k - 2]
            if item1 == item2:
                ret_list.append(item_set[i] | item_set[j])
    return ret_list


def apriori(dataset, min_support=0.5):
    """
    apriori 算法 主函数
    :param dataset: 数据集
    :param min_support: 最小支持度
    :return: 频繁项集, 相应支持度
    """
    # 创建初始候选集
    c1 = create_c1(dataset)
    # 扫描得出初始频繁项集
    l1, support_data = scan_d(dataset, c1, min_support=min_support)
    # 循环寻找多项频繁项集
    splice_set = [l1]
    k = 2   # 判断是否拼接所用
    while True:
        # 进行项的拼接
        ck = apriori_gen(splice_set[k - 2], k)
        # 扫描得出各项数频繁项集
        lk, support = scan_d(dataset, ck, min_support=min_support)
        # 如果包含频繁项集，继续循环，否则退出循环
        if len(ck) > 0:
            support_data.update(support)    # 字典更新
            splice_set.append(lk)
        else:
            break
        k += 1

    return splice_set, support_data


def generate_rules(splice_set, support_data, min_conf=0.6):
    """
    规则计算主函数
    :param splice_set: 频繁项集
    :param support_data: 相应支持度
    :param min_conf: 最小置信度
    :return: 规则列表
    """
    rule_list = []
    for i in range(1, len(splice_set)):
        for freq_set in splice_set[i]:
            h = [frozenset([item]) for item in freq_set]
            # 寻找满足置信度的规则并保存到 rule_list 中
            rules_from_conf(freq_set, h, support_data, rule_list, min_conf)
    return rule_list


def rules_from_conf(freq_set, h, support_data, rule_list, min_conf):
    """
    循环寻找各种规则
    freq_set: 项集
    h: 项集的中间产物，或者说是 Y ，而 freq_set - con_seq 就是 X 。
    rule_list: 满足条件的规则存储位置
    """
    m = len(h[0])
    while len(freq_set) > m:
        # cal_conf()修剪掉不满足的规则，保留满足规则
        h = cal_conf(freq_set, h, support_data, rule_list, min_conf)
        if len(h) > 1:
            # 当 项数 > 3 时，进行拼接 反向 和 深层 寻找，如：
            # 反向: {1, 2, 3} 对应 {1} ==> {2, 3} 等
            # 深层: {1, 2, 3, 4} 对应 {1, 2} ==> {3, 4} 等
            h = apriori_gen(h, m + 1)
            m += 1
        else:
            break


def cal_conf(freq_set, h, support_data, rule_list, min_conf):
    """
    对寻找到的规则进行最小置信度的计算与判断
    :return: 修剪后的项集
    """
    pruned_h = []
    # 循环遍历各项
    for con_seq in h:
        # 计算置信度 P(X==>Y) = P(XY) / P(X)
        conf = round(support_data[freq_set] / support_data[freq_set - con_seq], 2)
        # 判断是否满足最小置信度
        if conf >= min_conf:
            # 保存，set使得输出更加美观
            rule_list.append((set(freq_set - con_seq), set(con_seq), conf))
            pruned_h.append(con_seq)
    return pruned_h


if __name__ == '__main__':
    data = load_dataset()
    splice_set, support_data = apriori(data)
    for i in range(len(splice_set)):
        print('项集', i + 1, ':', )
        for item in splice_set[i]:
            print(f' {set(item)}', end='')
        print()
    print()
    rules = generate_rules(splice_set, support_data)
    for x, y, c in rules:
        print(x, '==>', y, 'conf:', c)
